from rest_framework_simplejwt.serializers import TokenObtainPairSerializer, TokenRefreshSerializer

from account.models import User
from account.serializers import UserSerializer
from account.services import AccountService


class MyTokenObtainPairSerializer(TokenObtainPairSerializer):
    @classmethod
    def get_token(cls, user):
        token = super().get_token(user)
        account_service = AccountService(user=user)

        # Add custom claims
        token['username'] = user.username
        token['permissions'] = account_service.get_dict_type_permissions()
        token['groups'] = user.get_user_groups()
        # ...

        return token


class CustomTokenRefreshSerializer(TokenRefreshSerializer):
    def validate(self, attrs):
        data = super().validate(attrs)
        refresh = self.token_class(attrs['refresh'])
        user_id = refresh['user_id']
        user = User.objects.get(id=user_id)
        account_service = AccountService(user=user)
        data.update(
            permissions=account_service.get_dict_type_permissions(),
            groups=user.get_user_groups(),
            userinfo=UserSerializer(user, many=False).data,
        )
        return data
